from typing import List, Dict, Any

from src.basedata.infrastructure.client.redis.repository import RedisRepository


class MongoDBRepository:
    """
    motor repository impl
    """

    def __init__(self, session, table_name):
        self.session = session
        self.table_name = table_name

    async def aggregate(self, pipeline: List[Dict[str, Any]], *args, **kwargs) -> List[Dict[str, Any]]:
        table = self.session.get_collection(self.table_name)
        return [record async for record in table.aggregate(pipeline, *args, **kwargs)]

    async def query_record(self, query: dict, *args, **kwargs):
        table = self.session.get_collection(self.table_name)
        return await table.find_one(query, *args, **kwargs)

    async def query_records(self, query: dict, *args, **kwargs):
        table = self.session.get_collection(self.table_name)
        cursor = table.find(query, *args, **kwargs)
        return [record async for record in cursor]

    async def query_records_limit(self, query: dict, page: int, num: int, sort={"updated_at": -1}):
        page = int(page)
        num = int(num)
        skip_num = num * (page - 1)
        table = self.session.get_collection(self.table_name)
        if sort:
            sort_field = [(key, value) for key, value in sort.items()]
            cursor = table.find(query).sort(sort_field).skip(skip_num).limit(num)
        else:
            cursor = table.find(query).skip(skip_num).limit(num)
        return [record async for record in cursor]

    async def query_count(self, query: dict):
        table = self.session.get_collection(self.table_name)
        return await table.count_documents(query)

    async def insert(self, record: dict):
        table = self.session.get_collection(self.table_name)
        result = await table.insert_one(record)
        return result.inserted_id

    async def insert_many(self, records: list) -> List[str]:
        table = self.session.get_collection(self.table_name)
        return await table.insert_many(records)

    async def update_one(self, query: dict, value: dict):
        table = self.session.get_collection(self.table_name)
        return await table.update_one(query, {'$set': value})

    async def update_many(self, query: dict, value: dict):
        table = self.session.get_collection(self.table_name)
        return await table.update_many(query, {'$set': value})

    async def delete(self, query: dict):
        table = self.session.get_collection(self.table_name)
        return await table.delete_one(query)

    async def bulk_write(self, data, ordered=False):
        table = self.session.get_collection(self.table_name)
        return await table.bulk_write(data, ordered)


class MongoDBMixedRedisRepository(MongoDBRepository, RedisRepository):
    def __init__(self, session, table_name, *args):
        super(MongoDBMixedRedisRepository, self).__init__(session, table_name)
        RedisRepository.__init__(self, *args)